import torch
import game_layer_cpp

device = 'cuda'
BATCH_SIZE = 1
BOARD_SIZE = 11
board = torch.zeros((BATCH_SIZE, 2, BOARD_SIZE, BOARD_SIZE), device=device, dtype=torch.bool)

#define STATUS_WIN 0
#define STATUS_OPEN4 1
#define STATUS_DEAD4 2
#define STATUS_OPEN3 3
#define STATUS_SLEEP3 4
#define STATUS_OPEN2 5

print(game_layer_cpp.get_status(board).shape)
print(game_layer_cpp.get_status(board).nonzero())

board[0][1][0][0] = 1
board[0][1][0][1] = 1
print("\nbelow should [4, 0, 2]")
print(game_layer_cpp.get_status(board).nonzero())

board[0][1][0][2] = 1
print("\nbelow should [2, 0, 3]")
print(game_layer_cpp.get_status(board).nonzero())

board[0][1][0][3] = 1
print("\nbelow should [0, 0, 4]")
print(game_layer_cpp.get_status(board).nonzero())

board.zero_()

board[0][1][0][1] = 1
print("\nbelow should [5, 0, 2]")
print(game_layer_cpp.get_status(board).nonzero())

board[0][1][0][2] = 1
print("\nbelow should [3, 0, 3], [4, 0, 0]")
print(game_layer_cpp.get_status(board).nonzero())

board[0][1][0][3] = 1
print("\nbelow should [1, 0, 4], [2, 0, 0]")
print(game_layer_cpp.get_status(board).nonzero())

board[0][1][0][4] = 1
print("\nbelow should [0, 0, 0], [0, 0, 5]")
print(game_layer_cpp.get_status(board).nonzero())


board = torch.rand((256, 2, 15, 15), device=device) < 0.1
print(game_layer_cpp.get_status(board).nonzero())
torch.cuda.synchronize()
